import shelve
import numpy as np
import json


class EasySample:
    """ "
    通过MD5值快速得到各个模型所需要的样本输入
    """

    sample = None

    def __init__(self) -> None:
        with open("config.json", "r") as file:
            self.config = json.load(file)
        self.idaDataPath = self.config["featureExtract"]["idaFeatureSaveDirectory"]

    def queryProcessSample(self, name, model="HGMSim"):
        """
        获得模型处理好的数据
        """
        if model == "HGMSim":
            return self.queryHGMSimSample(name)

    def queryHGMSimSample(self, name):
        """
        输出：
            adj 邻接矩阵 b*n*n
            att 特征矩阵 b*n*d
            vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot = self.idaDataPath + "//" + name
        res = {}
        res["adj"], res["att"], res["vtype"] = [], [], []
        with shelve.open(fileroot) as file:
            cg = file["cg"]
            cgattr = file["cgattr"]
            try:
                functype = file["funcType"]
            except Exception as e:
                functype = file["functype"]
            try:
                funcs_id = file["func_id"]  # name-->ind
            except Exception as e:
                funcs_id = file["funcs_id"]  # name-->ind
        id_funcs = {}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]] = i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData = []
        for i in functype.keys():
            temp = []
            if functype[i] == "local":
                temp = [1, 0, 0]
            elif functype[i] == "dynamic import":
                temp = [0, 1, 0]
                cgattr[i] = self.transApiName2Vector(id_funcs[i], 8, 20)
            else:
                temp = [0, 0, 1]
            tempData.append(temp)
        res["adj"] = np.array(cg)
        res["att"] = np.array(cgattr)
        res["vtype"] = np.array(tempData)
        return res

    def transApiName2Vector(self, function_name, embeddingSize, numLim):
        """
        将api转换成一个embeddingSize位的向量，每一位使用numLim取余
        """
        function_bytes = function_name.encode("utf-8")
        eight_bit_vector = [0] * embeddingSize
        for i, byte in enumerate(function_bytes):
            index = i % embeddingSize
            eight_bit_vector[index] += byte
            eight_bit_vector[index] %= numLim
        return eight_bit_vector


if __name__ == "__main__":
    a = EasySample()
    b = a.queryProcessSample("6e56ab5d8d53a5c8666007530f692f40")
    print(b)
